import numpy as np
from . import LocalSearchOperator, register_operator


class OrOptOperator(LocalSearchOperator):
    """Or-Opt局部搜索算子，通过移动路径中的连续段来改进解"""

    def __init__(self):
        self.name = "or_opt"
        self.supported_problems = ["tsp"]
        self.segment_lengths = [1, 2, 3]  # 尝试移动的段长度

    def generate_neighborhood(self, solution, problem_instance):
        """生成Or-Opt邻域"""
        neighborhood = []
        n = len(solution)

        # 尝试不同的段长度
        for segment_length in self.segment_lengths:
            # 遍历所有可能的段起始位置
            for i in range(1, n - segment_length):
                segment = solution[i:i + segment_length]

                # 生成移除段后的路径
                remaining_path = np.concatenate([solution[:i], solution[i + segment_length:]])

                # 尝试将段插入到所有可能的位置
                for j in range(1, len(remaining_path)):
                    if j == i:  # 跳过原始位置
                        continue

                    # 创建新解：将段插入到位置j
                    new_solution = np.concatenate([
                        remaining_path[:j],
                        segment,
                        remaining_path[j:]
                    ])

                    neighborhood.append(new_solution)

        return neighborhood

    def efficient_or_opt(self, solution, problem_instance, intensity=5):
        """更高效的Or-Opt实现，直接计算成本变化而不生成完整邻域"""
        n = len(solution)
        current_cost = problem_instance.evaluate(solution)
        best_solution = solution.copy()
        best_cost = current_cost
        improved = True
        iterations = 0

        while improved and iterations < intensity:
            improved = False

            for segment_length in self.segment_lengths:
                for i in range(1, n - segment_length):
                    # 计算移除段前后的边
                    prev_node = solution[i - 1]
                    segment_start = solution[i]
                    segment_end = solution[i + segment_length - 1]
                    next_node = solution[i + segment_length] if i + segment_length < n else solution[0]

                    # 计算移除段节省的成本
                    removal_saving = (
                            problem_instance.distance_matrix[prev_node][segment_start] +
                            problem_instance.distance_matrix[segment_end][next_node]
                    )

                    # 尝试将段插入到所有可能的位置
                    for j in range(1, n - segment_length + 1):
                        if j == i or j == i - 1:  # 跳过无效位置
                            continue

                        # 计算插入点前后的节点
                        insert_prev = solution[j - 1]
                        insert_next = solution[j] if j < n else solution[0]

                        # 计算插入段增加的成本
                        insertion_cost = (
                                problem_instance.distance_matrix[insert_prev][segment_start] +
                                problem_instance.distance_matrix[segment_end][insert_next]
                        )

                        # 计算净成本变化
                        net_change = insertion_cost - removal_saving

                        # 如果成本降低，则应用移动
                        if net_change < 0:
                            # 创建新解
                            segment = solution[i:i + segment_length]
                            remaining_path = np.concatenate([solution[:i], solution[i + segment_length:]])
                            new_solution = np.concatenate([
                                remaining_path[:j],
                                segment,
                                remaining_path[j:]
                            ])

                            # 更新最佳解
                            best_solution = new_solution
                            best_cost = current_cost + net_change
                            improved = True
                            break  # 找到一个改进就跳出内层循环

                    if improved:
                        break  # 找到一个改进就跳出段长度循环

                if improved:
                    break  # 找到一个改进就跳出段起始位置循环

            iterations += 1
            if improved:
                solution = best_solution
                current_cost = best_cost

        improvement = problem_instance.evaluate(solution) - current_cost
        return best_solution, improvement

    def __call__(self, solution, problem_instance, intensity=5):
        """应用Or-Opt局部搜索"""
        if not self.is_applicable(problem_instance):
            return solution, 0

        return self.efficient_or_opt(solution, problem_instance, intensity)


register_operator("or_opt", OrOptOperator())